#include "CubemapTextureArray.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Utils/FileUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

CubemapTextureArray::Builder::Builder() noexcept = default;
CubemapTextureArray::Builder::~Builder() noexcept = default;
CubemapTextureArray::Builder::Builder(CubemapTextureArray::Builder const& rhs) noexcept = default;
CubemapTextureArray::Builder::Builder(CubemapTextureArray::Builder && rhs) noexcept = default;
CubemapTextureArray::Builder& CubemapTextureArray::Builder::operator=(CubemapTextureArray::Builder const& rhs) noexcept = default;
CubemapTextureArray::Builder& CubemapTextureArray::Builder::operator=(CubemapTextureArray::Builder && rhs) noexcept = default;

CubemapTextureArray::Builder& CubemapTextureArray::Builder::SetAutoMip() noexcept
{
	mImpl->mipLevels = UINT32_MAX;
	return *this;
}

CubemapTextureArray::Builder& CubemapTextureArray::Builder::SetMipLevels(uint32 mipLevels) noexcept
{
	mImpl->mipLevels = mipLevels;
	return *this;
}

CubemapTextureArray::Builder& CubemapTextureArray::Builder::SetFormat(VkFormat format) noexcept
{
	assert(GStbFormat[format] >= 0 && GStbFormat[format] <= 4);
	mImpl->format = format;
	return *this;
}

CubemapTextureArray::Builder& CubemapTextureArray::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

CubemapTextureArray::Builder& CubemapTextureArray::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}

CubemapTextureArray::Builder& CubemapTextureArray::Builder::AddPath(std::array<std::string, 6> path) noexcept
{
	mImpl->paths.push_back(path);
	return *this;
}

RefCountPtr<CubemapTextureArray> CubemapTextureArray::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<CubemapTextureArray>(NewObject<CubemapTextureArray>(*mImpl));
}

RefCountPtr<CubemapTextureArray> CubemapTextureArray::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	CubemapTextureArray* cubemapTextureArray = NewObject<CubemapTextureArray>(*mImpl);
	process->AddTexture(name, cubemapTextureArray);
	return RefCountPtr<CubemapTextureArray>(cubemapTextureArray);
}

void CubemapTextureArray::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveTexture(mName, this);
		}
		DeleteObject<CubemapTextureArray>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveTexture(mName, this);
			}
			DeleteObject<CubemapTextureArray>(this);
			});
	}
}

CubemapTextureArray::CubemapTextureArray(BuilderDetails & builderDetails)
{
	mIsMultiple = false;
	mIsRenderTarget = false;

	mPaths = builderDetails.paths;

	// This flag is required for cube map images
	mFlags = VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT;
	mImageType = VK_IMAGE_TYPE_2D;
	mFormat = builderDetails.format;
	mExtent = { 0, 0, 1 };
	mMipLevels = builderDetails.mipLevels;
	// Cube faces count as array layers in Vulkan
	mArrayLayers = static_cast<uint32>(mPaths.size() * 6);
	mSamples = VK_SAMPLE_COUNT_1_BIT;
	mTiling = VK_IMAGE_TILING_OPTIMAL;
	mUsage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// Check if the queueFamilys is same
		bool isDifferent = false;
		const uint32 standQueueFamily = GVulkanRHI->GetQueueFamily(mQueueFamilys[0]);
		for (uint32 i = 1; i < mQueueFamilys.size(); ++i)
		{
			if (GVulkanRHI->GetQueueFamily(mQueueFamilys[i]) != standQueueFamily)
				isDifferent = true;
		}

		if (isDifferent)
			mSharingMode = VK_SHARING_MODE_CONCURRENT; // multi queue
		else
			mSharingMode = VK_SHARING_MODE_EXCLUSIVE;
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mImageLayouts.resize(1, VK_IMAGE_LAYOUT_UNDEFINED);
	mMemoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;

	// Cube map necessary
	mViewType = VK_IMAGE_VIEW_TYPE_CUBE_ARRAY;
	mComponents = { VK_COMPONENT_SWIZZLE_R , VK_COMPONENT_SWIZZLE_G , VK_COMPONENT_SWIZZLE_B , VK_COMPONENT_SWIZZLE_A };
	// the mip need to set later, the layer is the same to arraylayer
	mSubresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, VK_REMAINING_MIP_LEVELS, 0 ,mArrayLayers };

	mImages.resize(1, VK_NULL_HANDLE);
	mImageMemorys.resize(1, VK_NULL_HANDLE);
	mImageViews.resize(1, VK_NULL_HANDLE);

	// Get the first one pic for the size

	std::vector<uint8> buffer;
	FileUtil::LoadTextureToArray(mPaths[0][0], mFormat, buffer, mExtent.width, mExtent.height);

	// if mMipLevels is auto mip
	const uint32 maxMipLevels = static_cast<uint32>(std::floor(std::log2(glm::max<uint32>(mExtent.width, mExtent.height)))) + 1;
	mMipLevels = mMipLevels == UINT32_MAX ? maxMipLevels : mMipLevels;
	// if mMipLevels > 1, it is nesessary to copy data from image to make mip, check the miplevel is not too big
	if (mMipLevels > 1)
	{
		assert(mMipLevels <= maxMipLevels);
		mUsage |= VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
	}
	// Necessary to be set
	mSubresourceRange.levelCount = mMipLevels;

	// Each picture size
	VkDeviceSize dataSize = static_cast<VkDeviceSize>(mExtent.width * mExtent.height * mExtent.depth * GVkFormatSize[mFormat]);

	// Create stagebuffer to write data
	VkBuffer stagingBuffer;
	VkDeviceMemory stagingMemory;

	GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, dataSize * mArrayLayers, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

	void* data;
	VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, dataSize * mArrayLayers, 0, &data));

	memcpy(data, buffer.data(), dataSize);

	for (uint32 i = 0; i < mPaths.size(); ++i)
	{
		for (uint32 j = 0; j < 6; ++j)
		{
			// Ignore the first one pic
			if (i == 0 && j == 0)
				continue;

			data = static_cast<void*>(static_cast<uint8*>(data) + dataSize);

			uint32 width = 0;
			uint32 height = 0;
			FileUtil::LoadTextureToArray(mPaths[i][j], mFormat, buffer, width, height);

			assert(width == mExtent.width);
			assert(height == mExtent.height);

			memcpy(data, buffer.data(), dataSize);
		}
	}

	vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

	// Create image
	GVulkanRHI->CreateImage(mImages[0], mImageMemorys[0], mFlags, mExtent, mUsage, mFormat, mMipLevels, mMemoryProperty, mSamples, mImageType, mImageLayouts[0], mArrayLayers, mTiling, mSharingMode, mQueueFamilys);

	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	// TransitionImageLayout for copy data, all mip and layer are transfer to dst
	GVulkanRHI->TransferLayout(commandBuffer, mImages[0], mImageLayouts[0], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
	mImageLayouts[0] = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;

	// Copy data for miplevel 0 layer
	std::vector<VkBufferImageCopy> bufferImageCopys;
	bufferImageCopys.resize(mArrayLayers, {});
	for (uint32 i = 0; i < mArrayLayers; ++i)
	{
		bufferImageCopys[i] = {};
		bufferImageCopys[i].bufferOffset = dataSize * i;
		bufferImageCopys[i].bufferRowLength = 0;
		bufferImageCopys[i].bufferImageHeight = 0;
		bufferImageCopys[i].imageOffset = { 0, 0, 0 };
		bufferImageCopys[i].imageExtent = mExtent;
		bufferImageCopys[i].imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		bufferImageCopys[i].imageSubresource.mipLevel = 0;
		bufferImageCopys[i].imageSubresource.baseArrayLayer = i;
		bufferImageCopys[i].imageSubresource.layerCount = 1;
	}

	vkCmdCopyBufferToImage(
		commandBuffer,
		stagingBuffer,
		mImages[0],
		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		mArrayLayers,
		bufferImageCopys.data()
	);

	if (mMipLevels > 1)
	{
		GVulkanRHI->GenerateMipmaps2D(commandBuffer, mImages[0], mExtent.width, mExtent.height, mFormat, mMipLevels, mArrayLayers);
	}
	else
	{
		GVulkanRHI->TransferLayout(commandBuffer, mImages[0], mImageLayouts[0], VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);
	}
	mImageLayouts[0] = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

	GRenderState->FlushSingleCommand(commandBuffer);

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);

	GVulkanRHI->CreateImageView(mImageViews[0], mImages[0], mFormat, mViewType, mSubresourceRange, mComponents);
}
